#!/usr/bin/python
# -*- coding: utf-8 -*-  

"""
@Project : hello 
@file : xx_schedule_job_uat_dev.py
@Author : sheen
@time : 2025/4/10 12:55
@func : xx_schedule_job
重跑工具脚本 - 精确恢复任务状态

待优化:不需要重跑的任务不需要恢复任务状态

"""
import logging
import time
from datetime import datetime, timedelta
from typing import Dict, List, Optional, Tuple, Any
from com.cn.for_cdc.common.cdc_conmons import cdc_mysql_uat_cons
from com.cn.for_cdc.common.log_helper import configure_logging
from mysql_demo import MySQLHelper

configure_logging(log_format="[%(asctime)s] %(message)s")


class JobRerunner:
    def __init__(self):
        self.db_cache: Dict[str, MySQLHelper] = {}
        # 存储原始状态: {(bu, job_id): (status, restart, cron)}
        self.original_states: Dict[Tuple[str, int], Tuple[int, int, str]] = {}
        # 存储是否需要重跑的标识: {(bu, job_id): flag}
        self.rerun_flag: Dict[Tuple[str, int], bool] = {}
        # 存储是否需要重跑的标识: {(bu, job_id): flag}
        self.rerun_time: Dict[Tuple[str, int], str] = {}

    def _get_db_connection(self, bu: str) -> MySQLHelper:
        """获取数据库连接，使用缓存避免重复创建"""
        if bu not in self.db_cache:
            db_config = cdc_mysql_uat_cons[f'{bu}_uat']
            logging.info(f"创建数据库连接: {db_config}")
            self.db_cache[bu] = MySQLHelper(**db_config)
        return self.db_cache[bu]

    @staticmethod
    def generate_cron_text(delay_seconds: int = 20) -> tuple[str, str]:
        """生成cron表达式"""
        future_time = datetime.now() + timedelta(seconds=delay_seconds) - timedelta(hours=8)
        rerun_time = future_time.strftime("%Y-%m-%d %H:%M:%S")
        return rerun_time, f"{future_time.second} {future_time.minute} {future_time.hour} * * ?"

    def _backup_original_state(self, db: MySQLHelper, bu: str, job_id: int) -> None:
        """备份任务原始状态"""
        backup_sql = (
            f"SELECT STATUS, RESTART, CRON FROM asw_etl_{bu}_job.XX_SCHEDULE_JOB "
            f"WHERE job_id = {job_id};"
        )
        result = db.query(backup_sql)
        if result:
            original_state = result[0]
            self.original_states[(bu, job_id)] = (
                original_state['STATUS'],
                original_state['RESTART'],
                original_state['CRON']
            )
            logging.info(f"✅已备份原始状态: bu:{bu}, job_id:{job_id}, 原始状态:STATUS={original_state['STATUS']}, "
                         f"RESTART={original_state['RESTART']}, CRON={original_state['CRON']}")

    def _restore_original_state(self, db: MySQLHelper, bu: str, job_id: int) -> bool:
        """精确恢复任务原始状态"""
        if (bu, job_id) not in self.original_states:
            logging.info(f"未找到原始状态: bu:{bu}, job_id:{job_id}")
            return False

        original_status, original_restart, original_cron = self.original_states[(bu, job_id)]

        # 检查当前状态是否已经恢复
        check_sql = (
            f"SELECT STATUS, RESTART, CRON FROM asw_etl_{bu}_job.XX_SCHEDULE_JOB "
            f"WHERE job_id = {job_id};"
        )
        current_state = db.query(check_sql)[0]

        if (current_state['STATUS'] == original_status and
                current_state['RESTART'] == original_restart and
                current_state['CRON'] == original_cron):
            logging.info(f"状态已是最新无需恢复: bu:{bu}, job_id:{job_id}")
            return True

        # 执行恢复
        restore_sql = (
            f"UPDATE asw_etl_{bu}_job.XX_SCHEDULE_JOB "
            f"SET STATUS={original_status}, RESTART={original_restart}, "
            f"CRON='{original_cron}' "
            f"WHERE job_id = {job_id};"
        )
        logging.info(f"开始恢复原始状态: bu:{bu}, job_id:{job_id}\nSQL:\n{restore_sql}")
        db.execute(restore_sql)

        logging.info(f"已恢复原始状态: bu:{bu}, job_id:{job_id}\n"
                     f"恢复前: STATUS={current_state['STATUS']}, RESTART={current_state['RESTART']}, CRON={current_state['CRON']}\n"
                     f"恢复后: STATUS={original_status}, RESTART={original_restart}, CRON={original_cron}")
        return True

    def rerun_job(self, db: MySQLHelper, bu: str, job_id: int, force: bool = False) -> bool:
        """重跑任务"""
        logging.info(f"尝试开始重跑任务: bu:{bu}, job_id:{job_id}".center(100, '-'))
        flag = self.rerun_flag.get((bu, job_id), True)
        if not force:
            flag = self.check_job_rerun_required(db, bu, job_id)
            self.rerun_flag[(bu, job_id)] = flag
        if flag or force:
            # 备份原始状态
            self._backup_original_state(db, bu, job_id)
            rerun_time, cron = self.generate_cron_text()
            self.rerun_time[(bu, job_id)] = rerun_time
            rerun_sql = (
                f"UPDATE asw_etl_{bu}_job.XX_SCHEDULE_JOB "
                f"SET STATUS=1, RESTART=1, CRON='{cron}' "
                f"WHERE job_id = {job_id} ; "
            )
            affected_rows = db.execute(rerun_sql)
            logging.info(f"✅任务已重跑: bu:{bu}, job_id:{job_id} SQL:\n{rerun_sql} 受影响的行数: {affected_rows}")
            logging.info(f"bu:{bu}, job_id:{job_id}任务重跑时间: {rerun_time}")

            self.restore_state(db, bu, job_id, rerun_time, force)
        else:
            logging.info(f"任务未重跑: bu:{bu}, job_id:{job_id}")

    def restore_state(self, db: MySQLHelper, bu: str, job_id: int, rerun_time: str, force: bool = False) -> bool:
        """恢复原始状态"""
        logging.info(f"✅检查任务是否被调起: bu:{bu}, job_id:{job_id}".center(100, '-'))
        sql = f'''SELECT JOB_ID,JOB_NAME,STATUS,CREATED_TIME from asw_etl_{bu}_job.XX_JOB_LOG WHERE JOB_ID = {job_id} and CREATED_TIME >='{rerun_time}' ;'''
        logging.info(sql)
        while True:
            result = db.query(sql)
            if result:
                logging.info(f"✅bu:{bu}的job_id:{job_id} 已被调起,恢复之前的状态。")
                rerun_result = db.query(sql)
                logging.info(rerun_result)
                # 尝试恢复原始状态
                self._restore_original_state(db, bu, job_id)
                break
            else:
                logging.info(f"bu:{bu} 的job_id:{job_id} 任务还未被调起...")
                time.sleep(10)

    def check_failed_jobs(self, bu: str) -> List[int]:
        """检查所有失败的任务"""
        check_sql = f"""
            SELECT job_id
            FROM (
                SELECT 
                    job_id,
                    status,
                    RANK() OVER (PARTITION BY job_id ORDER BY created_time DESC) AS rn
                FROM asw_etl_{bu}_job.XX_JOB_LOG
                WHERE DATE(created_time) >= CURDATE() - 1
            ) t
            WHERE t.rn = 1 AND t.status = 'F' ;
        """
        db = self._get_db_connection(bu)
        res = db.query(check_sql)
        return [item['job_id'] for item in res] if res else []

    def batch_rerun(self, bu: str, job_ids: List[int], force: bool = False) -> None:
        """批量重跑任务"""
        db = self._get_db_connection(bu)

        # 先批量设置重跑标识
        for job_id in job_ids:
            self.rerun_job(db, bu, job_id, force)

    def check_rerun_status(self, bu: str, job_ids: List[int]) -> None:
        """检查重跑状态"""
        logging.info(f"开始检查重跑状态: bu:{bu}, job_ids:{job_ids}".center(100, '-'))
        db = self._get_db_connection(bu)

        for job_id in job_ids:
            if not self.rerun_flag[(bu, job_id)]:
                logging.info(f"✅任务状态正常无需重跑,也无需检查重跑状态: bu:{bu}, job_id:{job_id}")
            else:
                try:
                    rerun_time = self.rerun_time[(bu, job_id)]
                    while True:
                        time.sleep(30)
                        check_sql = (
                            f"SELECT JOB_ID,STATUS,CREATED_TIME FROM asw_etl_{bu}_job.XX_JOB_LOG "
                            f"WHERE job_id={job_id} and created_time >= '{rerun_time}' "
                            f"ORDER BY created_time DESC LIMIT 1;"
                        )
                        logging.info(check_sql)
                        res = db.query(check_sql)
                        logging.info(f"任务bu:{bu}, job_id:{job_id}最新日志: res:{res}")
                        if res and res[0]['STATUS'] == 'S':
                            logging.info(f"✅任务重跑成功: bu:{bu}, job_id:{job_id}")
                            break
                        if res and res[0]['STATUS'] == 'F':
                            logging.info(f"任务重跑失败: bu:{bu}, job_id:{job_id}")
                            break
                        logging.info(f"检查重跑状态,等待重跑完成: bu:{bu}, job_id:{job_id}")
                except:
                    logging.info(f"✅任务状态正常无需重跑,也无需检查重跑状态: bu:{bu}, job_id:{job_id}")

    def run(self, tasks: Dict[str, List[int]], force: bool = False) -> None:
        """执行重跑任务"""
        for bu, job_ids in tasks.items():
            self.batch_rerun(bu, job_ids, force)

        # 监测重跑的任务
        for bu, job_ids in tasks.items():
            self.check_rerun_status(bu, job_ids)

    def check_job_rerun_required(self, db, bu, job_id):
        """检查任务是否需要重跑"""
        # 检查任务状态
        check_sql = (
            f"SELECT JOB_ID,STATUS,CREATED_TIME FROM asw_etl_{bu}_job.XX_JOB_LOG "
            f"WHERE job_id={job_id} ORDER BY created_time DESC LIMIT 1;"
        )
        check_res = db.query(check_sql)
        logging.info(f"✅检查上次任务的状态")
        logging.debug(f"检查上次任务的状态,SQl:\n{check_sql}")
        logging.info(f"检查结果: bu:{bu}, job_id:{job_id} res:{check_res}")
        if check_res and check_res[0]['STATUS'] == 'F':
            logging.info(f"任务失败需重跑: bu:{bu}, job_id:{job_id}")
            return True
        logging.info(f"✅任务状态正常无需重跑: bu:{bu}, job_id:{job_id}")
        return False


if __name__ == '__main__':
    runner = JobRerunner()
    tasks = {
        "tp": [70361],
        # "drlt": [1270]
    }

    # 执行重跑(会自动恢复原始状态)
    runner.run(tasks)
